// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//


#include <pollux/functions/prestosql/json/json_extractor.h>

#include <cctype>
#include <unordered_map>
#include <vector>
#include <melon/container/f14_map.h>
#include "boost/algorithm/string/trim.hpp"
#include "melon/string.h"
#include "melon/json.h"
#include <pollux/common/base/exceptions.h>
#include <pollux/functions/prestosql/json/json_path_tokenizer.h>

namespace kumo::pollux::functions {
    namespace {
        using JsonVector = std::vector<const melon::Dynamic *>;

        class JsonExtractor {
        public:
            // Use this method to get an instance of JsonExtractor given a json path.
            static JsonExtractor &getInstance(melon::StringPiece path) {
                // Pre-process
                auto trimedPath = melon::trimWhitespace(path).str();

                std::shared_ptr<JsonExtractor> op;
                if (kExtractorCache.count(trimedPath)) {
                    op = kExtractorCache.at(trimedPath);
                } else {
                    if (kExtractorCache.size() == kMaxCacheNum) {
                        // TODO: Blindly evict the first one, use better policy
                        kExtractorCache.erase(kExtractorCache.begin());
                    }
                    op = std::make_shared<JsonExtractor>(trimedPath);
                    kExtractorCache[trimedPath] = op;
                }
                return *op;
            }

            melon::Optional<melon::Dynamic> extract(const melon::Dynamic &json);

            // Shouldn't instantiate directly - use getInstance().
            explicit JsonExtractor(const std::string &path) {
                if (!tokenize(path)) {
                    POLLUX_USER_FAIL("Invalid JSON path: {}", path);
                }
            }

        private:
            bool tokenize(const std::string &path) {
                if (path.empty()) {
                    return false;
                }
                if (!kTokenizer.reset(path)) {
                    return false;
                }

                while (kTokenizer.hasNext()) {
                    if (auto token = kTokenizer.getNext()) {
                        auto &tk = token.value();
                        if (tk.selector == JsonPathTokenizer::Selector::WILDCARD) {
                            tokens_.emplace_back("*");
                        } else if (tk.selector == JsonPathTokenizer::Selector::RECURSIVE) {
                            // Not supported in this implementation of json_extract
                            tokens_.clear();
                            return false;
                        } else {
                            tokens_.emplace_back(tk.value);
                        }
                    } else {
                        tokens_.clear();
                        return false;
                    }
                }
                return true;
            }

            // Cache tokenize operations in JsonExtractor across invocations in the same
            // thread for the same JsonPath.
            thread_local static std::
            unordered_map<std::string, std::shared_ptr<JsonExtractor> >
            kExtractorCache;
            thread_local static JsonPathTokenizer kTokenizer;

            // Max extractor number in extractor cache
            static const uint32_t kMaxCacheNum{32};

            std::vector<std::string> tokens_;
        };

        thread_local std::unordered_map<std::string, std::shared_ptr<JsonExtractor> >
        JsonExtractor::kExtractorCache;
        thread_local JsonPathTokenizer JsonExtractor::kTokenizer;

        void extractObject(
            const melon::Dynamic *jsonObj,
            const std::string &key,
            JsonVector &ret) {
            auto val = jsonObj->get_ptr(key);
            if (val) {
                ret.push_back(val);
            }
        }

        void extractArray(
            const melon::Dynamic *jsonArray,
            const std::string &key,
            JsonVector &ret) {
            int64_t arrayLen = jsonArray->size();
            if (key == "*") {
                for (size_t i = 0; i < arrayLen; ++i) {
                    ret.push_back(jsonArray->get_ptr(i));
                }
            } else {
                auto rv = melon::tryTo<int64_t>(key);
                if (rv.hasValue()) {
                    auto idx = rv.value();
                    idx = idx >= 0 ? idx : arrayLen + idx;
                    if (idx >= 0 && idx < arrayLen) {
                        ret.push_back(jsonArray->get_ptr(idx));
                    }
                }
            }
        }

        melon::Optional<melon::Dynamic> JsonExtractor::extract(
            const melon::Dynamic &json) {
            JsonVector input;
            // Temporary extraction result holder, swap with input after
            // each iteration.
            JsonVector result;
            input.push_back(&json);

            for (auto &token: tokens_) {
                for (auto &jsonObj: input) {
                    if (jsonObj->isObject()) {
                        extractObject(jsonObj, token, result);
                    } else if (jsonObj->isArray()) {
                        extractArray(jsonObj, token, result);
                    }
                }
                if (result.empty()) {
                    return melon::none;
                }
                input.clear();
                result.swap(input);
            }

            auto len = input.size();
            if (0 == len) {
                return melon::none;
            } else if (1 == len) {
                return *input.front();
            } else {
                melon::Dynamic array = melon::Dynamic::array;
                for (auto obj: input) {
                    array.push_back(*obj);
                }
                return array;
            }
        }

        bool isScalarType(const melon::Optional<melon::Dynamic> &json) {
            return json.has_value() && !json->isObject() && !json->isArray() &&
                   !json->isNull();
        }
    } // namespace

    melon::Optional<melon::Dynamic> jsonExtract(
        melon::StringPiece json,
        melon::StringPiece path) {
        try {
            // If extractor fails to parse the path, this will throw a PolluxUserError,
            // and we want to let this exception bubble up to the client. We only catch
            // json parsing failures (in which cases we return melon::none instead of
            // throw).
            auto &extractor = JsonExtractor::getInstance(path);
            return extractor.extract(melon::parseJson(json));
        } catch (const melon::json::parse_error &) {
        } catch (const melon::ConversionError &) {
            // Folly might throw a conversion error while parsing the input json. In
            // this case, let it return null.
        }
        return melon::none;
    }

    melon::Optional<melon::Dynamic> jsonExtract(
        const std::string &json,
        const std::string &path) {
        return jsonExtract(melon::StringPiece(json), melon::StringPiece(path));
    }

    melon::Optional<melon::Dynamic> jsonExtract(
        const melon::Dynamic &json,
        melon::StringPiece path) {
        try {
            return JsonExtractor::getInstance(path).extract(json);
        } catch (const melon::json::parse_error &) {
        }
        return melon::none;
    }

    melon::Optional<std::string> jsonExtractScalar(
        melon::StringPiece json,
        melon::StringPiece path) {
        auto res = jsonExtract(json, path);
        // Not a scalar value
        if (isScalarType(res)) {
            if (res->isBool()) {
                return res->asBool() ? std::string{"true"} : std::string{"false"};
            } else {
                return res->asString();
            }
        }
        return melon::none;
    }

    melon::Optional<std::string> jsonExtractScalar(
        const std::string &json,
        const std::string &path) {
        melon::StringPiece jsonPiece{json};
        melon::StringPiece pathPiece{path};

        return jsonExtractScalar(jsonPiece, pathPiece);
    }
} // namespace kumo::pollux::functions
